package main

import (
	"io"
	"log"
	"net"
	"os"
)

const OP_UPLOAD = 1
const OP_DOWNLOAD = 0
const UPLOAD_DIR = "E:\\新建文件夹\\GoLangDemo\\Zero-learn\\Beginner\\9tcp-trans-file\\upload"
const SOURCE_DIR = "E:\\新建文件夹\\GoLangDemo\\Zero-learn\\Beginner\\9tcp-trans-file\\sourceFile"

func main() {
	log.SetFlags(log.Llongfile)
	listener, err := net.Listen("tcp", ":9090")
	if err != nil {
		log.Println(err)
		return
	}
	defer listener.Close()
	log.Println("listen on :", listener.Addr().String())

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println(err)
			return
		}
		log.Println("new conn: ", conn.RemoteAddr())
		go handle(conn)
	}
}

func handle(conn net.Conn) {
	defer func() {
		log.Println("close conn:", conn.RemoteAddr())
		conn.Close()
	}()
	fileInfoBytes := make([]byte, 1024)
	_, err := conn.Read(fileInfoBytes)
	if err != nil {
		log.Println(err)
		return
	}
	op := int(fileInfoBytes[0])
	fileNameLen := int(fileInfoBytes[1])
	fileName := string(fileInfoBytes[2 : 2+fileNameLen])
	if op == OP_UPLOAD {
		err := recvFile(conn, UPLOAD_DIR+fileName)
		if err != nil {
			log.Println(err)
			return
		}
	} else {
		err := sendFile(conn, SOURCE_DIR+fileName)
		if err != nil {
			log.Println(err)
			return
		}
	}
}

func recvFile(conn net.Conn, filePath string) (err error) {
	file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		log.Println(err)
		return
	}
	defer file.Close()
	buf := make([]byte, 512)
	n := 0
	for {
		n, err = conn.Read(buf)
		if err != nil && err != io.EOF {
			log.Println(err)
			return
		} else {
			err = nil
		}
		if n == 0 {
			break
		}
		file.Write(buf)
	}
	return
}

func sendFile(conn net.Conn, filePath string) (err error) {
	src, err := os.Open(filePath)
	if err != nil {
		log.Println(err)
		return
	}
	defer src.Close()
	buf := make([]byte, 1024)
	n := 0
	for {
		n, err = src.Read(buf)
		if err != nil && err != io.EOF {
			log.Println(err)
			return
		} else {
			err = nil
		}
		if n == 0 {
			break
		}
		if _, err = conn.Write(buf[:n]); err != nil {
			log.Println(err)
			return
		}
	}
	return
}
